#!/usr/bin/env python3
"""
DeepSeek AI故事生成系统 - 主程序入口
"""

import os
import sys
import asyncio
import yaml
from pathlib import Path
from loguru import logger
from dotenv import load_dotenv

# 添加项目根目录到Python路径
project_root = Path(__file__).parent
sys.path.insert(0, str(project_root))

from src.services.deepseek_service import DeepSeekService, DeepSeekConfig
from src.models.api_model import StoryRequest, StoryType, Language
from src.utils.database import DatabaseManager


class StoryGeneratorSystem:
    """故事生成系统主类"""
    
    def __init__(self):
        """初始化系统"""
        self.config = None
        self.deepseek_service = None
        self.database_manager = None
        self._load_config()
        self._init_services()
    
    def _load_config(self):
        """加载配置文件"""
        try:
            load_dotenv()
            config_path = project_root / "config" / "config.yaml"
            with open(config_path, 'r', encoding='utf-8') as f:
                self.config = yaml.safe_load(f)
            self._replace_env_vars()
            logger.info("配置文件加载成功")
        except Exception as e:
            logger.error(f"配置文件加载失败: {str(e)}")
            sys.exit(1)
    
    def _replace_env_vars(self):
        """替换环境变量"""
        def replace_in_dict(data):
            if isinstance(data, dict):
                for key, value in data.items():
                    if isinstance(value, str) and value.startswith("${") and value.endswith("}"):
                        env_var = value[2:-1]
                        data[key] = os.getenv(env_var, "")
                    elif isinstance(value, (dict, list)):
                        replace_in_dict(value)
            elif isinstance(data, list):
                for item in data:
                    replace_in_dict(item)
        replace_in_dict(self.config)
    
    def _init_services(self):
        """初始化服务"""
        try:
            # 初始化DeepSeek服务
            deepseek_config = DeepSeekConfig(
                api_key=self.config["api"]["deepseek"]["api_key"],
                base_url=self.config["api"]["deepseek"]["base_url"],
                model=self.config["api"]["deepseek"]["model"],
                max_tokens=self.config["api"]["deepseek"]["max_tokens"],
                temperature=self.config["api"]["deepseek"]["temperature"],
                timeout=self.config["api"]["deepseek"]["timeout"]
            )
            self.deepseek_service = DeepSeekService(deepseek_config)
            
            # 初始化数据库管理器
            db_path = self.config["database"]["sqlite"]["path"]
            self.database_manager = DatabaseManager(db_path)
            
            logger.info("服务初始化完成")
        except Exception as e:
            logger.error(f"服务初始化失败: {str(e)}")
            sys.exit(1)
    
    async def test_system(self):
        """测试系统功能"""
        logger.info("开始系统测试...")
        
        # 测试API连接
        if await self.deepseek_service.test_connection():
            logger.info("✅ DeepSeek API连接测试成功")
        else:
            logger.error("❌ DeepSeek API连接测试失败")
            return False
        
        # 测试故事生成
        test_request = StoryRequest(
            title="测试故事",
            story_type=StoryType.SHUIHU,
            language=Language.BILINGUAL,
            character_name="宋江",
            length_limit=1000,
            bilingual=True,
            temperature=0.7
        )
        
        try:
            response = await self.deepseek_service.generate_story(test_request)
            if response.success:
                logger.info("✅ 故事生成测试成功")
                return True
            else:
                logger.error(f"❌ 故事生成测试失败: {response.message}")
                return False
        except Exception as e:
            logger.error(f"❌ 故事生成测试异常: {str(e)}")
            return False
    
    async def generate_shuihu_stories(self, start_episode: int = 1, count: int = 3):
        """生成水浒故事系列"""
        logger.info(f"开始生成水浒故事，从第{start_episode}集开始，共{count}集")
        
        # 水浒人物列表
        shuihu_characters = [
            {"name": "宋江", "nickname": "及时雨"},
            {"name": "卢俊义", "nickname": "玉麒麟"},
            {"name": "吴用", "nickname": "智多星"},
            {"name": "公孙胜", "nickname": "入云龙"},
            {"name": "关胜", "nickname": "大刀"}
        ]
        
        requests = []
        for i in range(count):
            episode_num = start_episode + i
            if episode_num <= len(shuihu_characters):
                character = shuihu_characters[episode_num - 1]
                request = StoryRequest(
                    title=f"小艾说水浒_第{episode_num}集_{character['nickname']}{character['name']}",
                    story_type=StoryType.SHUIHU,
                    language=Language.BILINGUAL,
                    character_name=character['name'],
                    length_limit=3000,
                    bilingual=True,
                    temperature=0.7
                )
                requests.append(request)
        
        responses = await self.deepseek_service.batch_generate_stories(requests)
        success_count = sum(1 for r in responses if r.success)
        logger.info(f"故事生成完成: 成功{success_count}个，失败{len(responses)-success_count}个")
        return responses
    
    def save_stories_to_files(self, responses, folder_name: str = None):
        """将故事保存到文件"""
        if not folder_name:
            from datetime import datetime
            folder_name = f"{datetime.now().strftime('%Y%m%d')}_01"
        
        folder_path = project_root / folder_name
        folder_path.mkdir(exist_ok=True)
        
        for response in responses:
            if response.success and response.title and response.content_zh:
                content = f"# {response.title}\n\n"
                content += "## 中文版\n\n"
                content += response.content_zh + "\n\n"
                
                if response.content_en:
                    content += "---\n\n"
                    content += "## English Version\n\n"
                    content += response.content_en + "\n\n"
                
                filename = f"{response.title}.txt"
                file_path = folder_path / filename
                
                with open(file_path, 'w', encoding='utf-8') as f:
                    f.write(content)
                
                logger.info(f"故事已保存: {file_path}")
    
    async def run_demo(self):
        """运行演示程序"""
        logger.info("🚀 启动DeepSeek AI故事生成系统演示")
        
        if not await self.test_system():
            logger.error("系统测试失败，程序退出")
            return
        
        logger.info("开始生成水浒故事演示...")
        responses = await self.generate_shuihu_stories(start_episode=4, count=3)
        self.save_stories_to_files(responses)
        logger.info("🎉 演示程序运行完成！")


async def main():
    """主函数"""
    try:
        system = StoryGeneratorSystem()
        await system.run_demo()
    except KeyboardInterrupt:
        logger.info("程序被用户中断")
    except Exception as e:
        logger.error(f"程序运行异常: {str(e)}")
        sys.exit(1)


if __name__ == "__main__":
    asyncio.run(main()) 